#!/usr/bin/python2.7
import matplotlib
#matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.animation import FFMpegWriter
from matplotlib.patches import Polygon
from matplotlib.collections import PatchCollection
import numpy as np
import math
import time
import datetime
import random
import copy
import os
import gc
import csv
import pickle
from graph import *
from robots import *


#plt.ion()

class Aera:
    def __init__(self, areasize=10, gridres=1.0, randmod=True):
        self.gridres = gridres
        self.areasize = areasize
        self.fig=plt.figure(num=1,figsize=(12, 12))
        self.iter = 0
        self.grids=[]
        self.ocgrids=[]
        self.randmod=randmod
        self.redraw()
        self.treepoint=[]
        self.trjpoint=[]
        self.smoothtrj=[]
        self.sorttree={}


    def redraw(self,numr=1):
        self.fig=plt.figure(num=numr, figsize=(12, 12))
        plt.clf()
        self.ax = self.fig.add_subplot(111)
        plt.xlim((0,10))
        plt.ylim((0, 10))
        #plt.hold(True)
        if(numr<=1):
            self.grid()
        print "Drawing graph vertexes.. Totally %d" % len(self.graph.vertexes)
        for v in self.graph.vertexes:
            plt.plot(self.graph.vertexes[v].center.x,self.graph.vertexes[v].center.y,'ro')
            rect=matplotlib.patches.Rectangle((self.graph.vertexes[v].center.x-self.gridres/2,self.graph.vertexes[v].center.y-self.gridres/2),self.gridres/2,self.gridres/2,edgecolor='y',fill=False,linestyle=':')
            self.ax.add_patch(rect)
            rect=matplotlib.patches.Rectangle((self.graph.vertexes[v].center.x,self.graph.vertexes[v].center.y-self.gridres/2),self.gridres/2,self.gridres/2,edgecolor='y',fill=False,linestyle=':')
            self.ax.add_patch(rect)
            rect=matplotlib.patches.Rectangle((self.graph.vertexes[v].center.x-self.gridres/2,self.graph.vertexes[v].center.y),self.gridres/2,self.gridres/2,edgecolor='y',fill=False,linestyle=':')
            self.ax.add_patch(rect)
            rect=matplotlib.patches.Rectangle((self.graph.vertexes[v].center.x,self.graph.vertexes[v].center.y),self.gridres/2,self.gridres/2,edgecolor='y',fill=False,linestyle=':')
            self.ax.add_patch(rect)
        for v in self.ocgrids:
            rect = matplotlib.patches.Rectangle((v.x-0.5, v.y-0.5),
                                                self.gridres, self.gridres, edgecolor='gray', fill=True,
                                                facecolor='gray')
            self.ax.add_patch(rect)

        print "Graph vertexes have been drawn."
        gc.collect()
        return self.grids

    def drawbackground(self,fignum):
        fig=plt.figure(num=fignum, figsize=(12, 12))
        ax = fig.add_subplot(111)
        plt.xlim((0, 10))
        plt.ylim((0, 10))
        plt.hold(True)
        print "Drawing graph vertexes.. Totally %d" % len(self.graph.vertexes)
        for v in self.graph.vertexes:
            plt.plot(self.graph.vertexes[v].center.x, self.graph.vertexes[v].center.y, 'ro',ms=3)
            rect = matplotlib.patches.Rectangle((self.graph.vertexes[v].center.x - self.gridres / 2,
                                                 self.graph.vertexes[v].center.y - self.gridres / 2), self.gridres / 2,
                                                self.gridres / 2, edgecolor='y', fill=False, linestyle=':')
            ax.add_patch(rect)
            rect = matplotlib.patches.Rectangle(
                (self.graph.vertexes[v].center.x, self.graph.vertexes[v].center.y - self.gridres / 2), self.gridres / 2,
                self.gridres / 2, edgecolor='y', fill=False, linestyle=':')
            ax.add_patch(rect)
            rect = matplotlib.patches.Rectangle(
                (self.graph.vertexes[v].center.x - self.gridres / 2, self.graph.vertexes[v].center.y), self.gridres / 2,
                self.gridres / 2, edgecolor='y', fill=False, linestyle=':')
            ax.add_patch(rect)
            rect = matplotlib.patches.Rectangle((self.graph.vertexes[v].center.x, self.graph.vertexes[v].center.y),
                                                self.gridres / 2, self.gridres / 2, edgecolor='y', fill=False,
                                                linestyle=':')
            ax.add_patch(rect)
        print "There are %d ocgrids"%(len(self.ocgrids))
        for v in self.ocgrids:
            rect = matplotlib.patches.Rectangle((v.x - 0.5, v.y - 0.5),
                                                self.gridres, self.gridres, edgecolor='gray', fill=True,
                                                facecolor='gray')
            ax.add_patch(rect)
        gc.collect()
        return ax


    def grid(self):
        cc = 0
        k = 0
        f = 0
        for i in range(self.areasize):
            for j in range(self.areasize):
                # if not self.randmod:
                #     if j>4 and j<6 and i>4 and i+j<8:
                #         self.ocgrids.append(vector2d(i + 0.5, j + 0.5))
                #         continue
                #     if j==5 and i>1 and i<9:
                #         self.ocgrids.append(vector2d(i + 0.5, j + 0.5))
                #         continue

                #     if  j<6 and i<5 and i>3:
                #         self.ocgrids.append(vector2d(i + 0.5, j + 0.5))
                #         continue

                #     if  j<3 and i>7:
                #         self.ocgrids.append(vector2d(i + 0.5, j + 0.5))
                #         continue

                #     if  j<2 and i<1:
                #         self.ocgrids.append(vector2d(i + 0.5, j + 0.5))
                #         continue
                # else:
                #     if f < 1 and random.random() > 0.95:
                #         cc = random.randint(3,self.areasize-5)
                #         k = 1
                #         f = 1.5*self.areasize
                #     if k < cc:
                #         k = k+1
                #         self.ocgrids.append(vector2d(i + 0.5, j + 0.5))
                #         continue
                #     else:
                #         f = f-1
                self.grids.append(vector2d(i+0.5,j+0.5))
        print "Creating graph vertexes.."
        self.graph = stcgraph(self.grids)
        print "Graph vertexes has been created."


    def spantree(self):
        v=vector2d(0.25,0.25)
        self.graph.span(v)
        plt.pause(0.5)
        for v in self.graph.graphvertexes:
            for e in self.graph.graphvertexes[v].edges:
                plt.plot([self.graph.graphvertexes[v].center.x,self.graph.graphvertexes[v].edges[e].center.x],[self.graph.graphvertexes[v].center.y,self.graph.graphvertexes[v].edges[e].center.y],'k')

        trjt=self.graph.gentrj(vector2d(0.25,0.25),self.graph.graphvertexes[str(0.5*1000+0.5)])
        trj=np.asarray(trjt)
        for i in range(len(trj)-1):
            plt.plot([trj[i,0],trj[i+1,0]],[trj[i,1],trj[i+1,1]],'g')
            plt.pause(0.1)


    def initani(self):
        print "initialize"
        tpi=self.treepoint[0]
        self.treeplot=[]
        self.trjplot=[]
        self.agentplot=[]
        for i in range(len(tpi[0])):
            # t,=self.plot([tpi[0][i].x,tpi[1][i].x],[tpi[0][i].y,tpi[1][i].y],color=(0,1-l*1.0/len(lp),l*1.0/len(lp)),lw=2)
            # self.treeplot.append(t)
            t, =self.ax.plot(self.trjpoint[i][0, 0], self.trjpoint[i][0, 1], ':',
                     color=(1 - i * 1.0 / len(self.trjpoint), i * 1.0 / len(self.trjpoint), i * 1.0 / len(self.trjpoint)))
            self.trjplot.append(t)
            t, =self.ax.plot(self.trjpoint[i][0, 0], self.trjpoint[i][0, 1], 'o', ms=10,
                     color=(1 - i * 1.0 / len(self.trjpoint), i * 1.0 / len(self.trjpoint), i * 1.0 / len(self.trjpoint)))
            self.agentplot.append(t)
        return tuple(self.trjplot)+tuple(self.agentplot)


    def updateani(self,frames):
        ltr = len(self.treepoint)
        frame=frames-5
        if frame<0:
            return tuple(self.trjplot) + tuple(self.agentplot)
        if frame < ltr:
            lp=self.treepoint[frame][0]
            # for i in range(len(lp)):
            #     self.treeplot.set_data([t.x for t in self.treepoint[0]])
            for i in range(len(lp)):
                self.ax.plot([self.treepoint[frame][0][i].x, self.treepoint[frame][1][i].x],
                         [self.treepoint[frame][0][i].y, self.treepoint[frame][1][i].y],
                          color=(0, 1 - i * 1.0 / len(lp), i * 1.0 / len(lp)), lw=2)
                self.ax.plot(self.treepoint[frame][0][i].x,
                         self.treepoint[frame][0][i].y,
                         'bo')
                self.ax.plot(self.treepoint[frame][1][i].x,
                         self.treepoint[frame][1][i].y,
                         'ro')
                if frame > 0:
                    plt.plot(self.treepoint[frame-1][1][i].x,
                            self.treepoint[frame-1][1][i].y,
                            'bo')
        else:
            pass
            for i in range(len(self.trjpoint)):
                try:
                    self.agentplot[i].set_data(self.trjpoint[i][frame-ltr-1,0],self.trjpoint[i][frame-ltr-1,1])
                    self.trjplot[i].set_data(self.trjpoint[i][0:frame - ltr, 0], self.trjpoint[i][0:frame - ltr, 1])
                except:
                    pass
        return tuple(self.trjplot)+tuple(self.agentplot)

    def updateani2(self,frames):
        ltr = len(self.sorttree[0])
        lp = len(self.sorttree)
        frame=frames-5
        if frame<1:
            return tuple(self.trjplot) + tuple(self.agentplot)
        if frame < ltr:
            # for i in range(len(lp)):
            #     self.treeplot.set_data([t.x for t in self.treepoint[0]])
            for i in range(len(self.sorttree)):
                if(len(self.sorttree[i])<=frame):
                    continue
                self.ax.plot(self.sorttree[i][frame].center.x,
                             self.sorttree[i][frame].center.y,
                         'bo')
                for e in self.sorttree[i][frame].edges:
                    self.ax.plot([self.sorttree[i][frame].center.x, self.sorttree[i][frame].edges[e].center.x],
                                 [self.sorttree[i][frame].center.y,  self.sorttree[i][frame].edges[e].center.y],
                                 color=(0, 1 - i * 1.0 / lp, i * 1.0 / lp), lw=2)
        else:
            pass
            for i in range(len(self.trjpoint)):
                try:
                    self.agentplot[i].set_data(self.trjpoint[i][frame-ltr-1,0],self.trjpoint[i][frame-ltr-1,1])
                    self.trjplot[i].set_data(self.trjpoint[i][0:frame - ltr, 0], self.trjpoint[i][0:frame - ltr, 1])
                except:
                    pass
        return tuple(self.trjplot)+tuple(self.agentplot)

    def animate(self,name="000"):
        saveit=True
        print len(self.treepoint)+len(self.trjpoint[0])
        anim = animation.FuncAnimation(self.fig,self.updateani,frames=range(len(self.treepoint)+len(self.trjpoint[0])+5),
                                       init_func=self.initani,interval=200,blit=False,repeat=False)
        print "animate started"
        if saveit:
            plt.ion()
        else:
            plt.show()
        if saveit:
            metadata = dict(title='Swarm Movie Test', artist='Dr. Dong', comment='Movie support!')
            writerf = FFMpegWriter(fps=15, metadata=metadata)
            anim.save(name+'.mp4', writer=writerf)
            self.saveeps(name)

    def animate2(self,name="000"):
        self.redraw(100)
        saveit=True
        print len(self.treepoint)+len(self.trjpoint[0])
        anim = animation.FuncAnimation(self.fig,self.updateani2,frames=range(len(self.treepoint)+len(self.trjpoint[0])+5),
                                       init_func=self.initani,interval=200,blit=False,repeat=False)
        print "animate started"
        if saveit:
            plt.ion()
        else:
            plt.show()
        if saveit:
            metadata = dict(title='Swarm Movie Test', artist='Dr. Dong', comment='Movie support!')
            writerf = FFMpegWriter(fps=15, metadata=metadata)
            anim.save(name+'.mp4', writer=writerf)
            self.saveeps2(name)


    def saveeps(self,name='000'):
        ax=self.drawbackground(2)
        for frame in range(len(self.treepoint)):
            lp = self.treepoint[frame][0]
            for i in range(len(lp)):
                ax.plot([self.treepoint[frame][0][i].x, self.treepoint[frame][1][i].x],
                         [self.treepoint[frame][0][i].y, self.treepoint[frame][1][i].y],
                         color=(0, 1 - i * 1.0 / len(lp), i * 1.0 / len(lp)), lw=2)
                ax.plot(self.trjpoint[i][0, 0], self.trjpoint[i][0, 1], 'o', ms=8,
                         color=(1 - i * 1.0 / len(self.trjpoint), i * 1.0 / len(self.trjpoint),
                                i * 1.0 / len(self.trjpoint)))
                ax.plot(self.trjpoint[i][:, 0], self.trjpoint[i][:, 1], ':',
                         color=(1 - i * 1.0 / len(self.trjpoint), i * 1.0 / len(self.trjpoint),
                                i * 1.0 / len(self.trjpoint)))
                ax.arrow(self.trjpoint[i][0, 0], self.trjpoint[i][0, 1],
                         (self.trjpoint[i][1, 0] - self.trjpoint[i][0, 0]) * 0.4,
                         (self.trjpoint[i][1, 1] - self.trjpoint[i][0, 1]) * 0.4, head_width=0.1,
                         color=(1 - i * 1.0 / len(self.trjpoint), i * 1.0 / len(self.trjpoint),
                                i * 1.0 / len(self.trjpoint)))
        plt.savefig(name + ".eps")

        for k in range(len(self.treepoint[0][0])):
            ax=self.drawbackground(k + 3)
            for frame in range(len(self.treepoint)):
                lp = self.treepoint[frame][0]
                for i in range(len(lp)):
                    if not k == i:
                        continue
                    ax.plot([self.treepoint[frame][0][i].x, self.treepoint[frame][1][i].x],
                             [self.treepoint[frame][0][i].y, self.treepoint[frame][1][i].y],
                             color=(0, 1 - i * 1.0 / len(lp), i * 1.0 / len(lp)), lw=2)
                    ax.plot(self.trjpoint[i][0, 0], self.trjpoint[i][0, 1], 'o', ms=8,
                             color=(1 - i * 1.0 / len(self.trjpoint), i * 1.0 / len(self.trjpoint),
                                    i * 1.0 / len(self.trjpoint)))
                    ax.plot(self.trjpoint[i][:, 0], self.trjpoint[i][:, 1], '--',
                             color=(1 - i * 1.0 / len(self.trjpoint), i * 1.0 / len(self.trjpoint),
                                    i * 1.0 / len(self.trjpoint)))
                    ax.arrow(self.trjpoint[i][0, 0], self.trjpoint[i][0, 1],
                             (self.trjpoint[i][1, 0] - self.trjpoint[i][0, 0]) * 0.4,
                             (self.trjpoint[i][1, 1] - self.trjpoint[i][0, 1]) * 0.4, head_width=0.1,
                             color=(1 - i * 1.0 / len(self.trjpoint), i * 1.0 / len(self.trjpoint),
                                    i * 1.0 / len(self.trjpoint)))

            ax.plot([x[0] for x in self.smoothtrj[k]],[x[1] for x in self.smoothtrj[k]],':',color=(1 - i * 1.0 / len(self.trjpoint), i * 1.0 / len(self.trjpoint),
                                    i * 1.0 / len(self.trjpoint)),lw=1)
            plt.savefig(name + "_%d.eps" % (k))


    def saveeps2(self, name='000'):
            lp=len(self.sorttree)
            ax = self.drawbackground(5+lp)
            for frame in range(len(self.sorttree[0])):
                for i in range(len(self.sorttree)):
                    if (len(self.sorttree[i]) <= frame):
                        continue
                    for e in self.sorttree[i][frame].edges:
                        ax.plot([self.sorttree[i][frame].center.x, self.sorttree[i][frame].edges[e].center.x],
                                [self.sorttree[i][frame].center.y, self.sorttree[i][frame].edges[e].center.y],
                                color=(0, 1 - i * 1.0 / lp, i * 1.0 / lp), lw=2)
                    ax.plot(self.trjpoint[i][0, 0], self.trjpoint[i][0, 1], 'o', ms=8,
                            color=(1 - i * 1.0 / len(self.trjpoint), i * 1.0 / len(self.trjpoint),
                                   i * 1.0 / len(self.trjpoint)))
                    ax.plot(self.trjpoint[i][:, 0], self.trjpoint[i][:, 1], ':',
                            color=(1 - i * 1.0 / len(self.trjpoint), i * 1.0 / len(self.trjpoint),
                                   i * 1.0 / len(self.trjpoint)))
                    ax.arrow(self.trjpoint[i][0, 0], self.trjpoint[i][0, 1],
                             (self.trjpoint[i][1, 0] - self.trjpoint[i][0, 0]) * 0.4,
                             (self.trjpoint[i][1, 1] - self.trjpoint[i][0, 1]) * 0.4, head_width=0.1,
                             color=(1 - i * 1.0 / len(self.trjpoint), i * 1.0 / len(self.trjpoint),
                                    i * 1.0 / len(self.trjpoint)))
            plt.savefig(name + ".eps")

            for k in range(len(self.sorttree)):
                ax = self.drawbackground(k + 6+lp)
                for frame in range(len(self.sorttree[0])):
                    for i in range(len(self.sorttree)):
                        if (len(self.sorttree[i]) <= frame):
                            continue
                        if not k == i:
                            continue
                        for e in self.sorttree[i][frame].edges:
                            ax.plot([self.sorttree[i][frame].center.x, self.sorttree[i][frame].edges[e].center.x],
                                    [self.sorttree[i][frame].center.y, self.sorttree[i][frame].edges[e].center.y],
                                    color=(0, 1 - i * 1.0 / lp, i * 1.0 / lp), lw=2)
                        ax.plot(self.trjpoint[i][0, 0], self.trjpoint[i][0, 1], 'o', ms=8,
                                color=(1 - i * 1.0 / len(self.trjpoint), i * 1.0 / len(self.trjpoint),
                                       i * 1.0 / len(self.trjpoint)))
                        ax.plot(self.trjpoint[i][:, 0], self.trjpoint[i][:, 1], '--',
                                color=(1 - i * 1.0 / len(self.trjpoint), i * 1.0 / len(self.trjpoint),
                                       i * 1.0 / len(self.trjpoint)))
                        ax.arrow(self.trjpoint[i][0, 0], self.trjpoint[i][0, 1],
                                 (self.trjpoint[i][1, 0] - self.trjpoint[i][0, 0]) * 0.4,
                                 (self.trjpoint[i][1, 1] - self.trjpoint[i][0, 1]) * 0.4, head_width=0.1,
                                 color=(1 - i * 1.0 / len(self.trjpoint), i * 1.0 / len(self.trjpoint),
                                        i * 1.0 / len(self.trjpoint)))

                ax.plot([x[0] for x in self.smoothtrj[k]],[x[1] for x in self.smoothtrj[k]],':',color=(1 - i * 1.0 / len(self.trjpoint), i * 1.0 / len(self.trjpoint),
                                    i * 1.0 / len(self.trjpoint)),lw=1)
                plt.savefig(name + "_%d.eps" % (k))
            plt.figure(500,figsize=(16,3.5))
            fig,ax=plt.subplots(1,2,figsize=(16,4))
            for k in range(2):
                t=np.arange(0,20,0.01)
                ax[k].plot(t,[x[0] for x in self.smoothtrj[k][0:2000]],'r--')
                ax[k].plot(t,[x[1] for x in self.smoothtrj[k][0:2000]])
                ax[k].grid()
                ax[k].legend(['X','Y'])
                ax[k].set_xlabel('Time (s) R%d'%k)
                if(k<1):
                    ax[k].set_ylabel('Position (m)')
            plt.savefig(name + "s.eps")


if __name__ == "__main__":
    rmod=False
    a = Aera(randmod=rmod)
    instance.copygraph(copy.deepcopy(a.graph))
    ags = []
    for i in range(3):
        aa = agent()
        ags.append(aa)

    ipos=instance.initagents(randmod=rmod)
    plt.plot(ipos[:,0],ipos[:,1],'k*',ms=10)
    tp=[]
    for j in range(100):
        lp, pp = instance.spanonce()
        if(len(pp)<1):
            break
        tpi=[]
        tpi.append(lp)
        tpi.append(pp)
        tp.append(tpi)


        # for l in range(len(lp)):
        #     plt.plot([lp[l].x, pp[l].x], [lp[l].y, pp[l].y],color=(0,1-l*1.0/len(lp),l*1.0/len(lp)),lw=2)
        #     plt.plot(lp[l].x, lp[l].y, 'o',color=(0,1-l*1.0/len(lp),l*1.0/len(lp)))
        #     plt.plot(pp[l].x, pp[l].y, 'bo')
            #plt.pause(0.1)
            #plt.plot(pp[l].x, pp[l].y, 'o',color=(0,1-l*1.0/len(lp),l*1.0/len(lp)))
    t=datetime.datetime.now()
    
    trj,smoothtrj= instance.gentrj()
    j=0
    # for t in trj:
    #     plt.arrow(t[0, 0], t[0, 1],(t[1,0]-t[0,0])*0.2,(t[1,1]-t[0,1])*0.2,head_width=0.1)

    a.treepoint=tp
    a.trjpoint=trj
    a.smoothtrj=smoothtrj
    
    #a.animate("./img/%02d%02d%02d%02d%02da"%(t.month,t.day,t.hour,t.minute,t.second))
    #a.saveeps("./img/%02d%02d%02d%02d%02db" % (t.month, t.day, t.hour, t.minute, t.second))
    
    a.sorttree = instance.sorttree()
    trj,smoothtrj = instance.gentrj()
    a.trjpoint = trj
    a.smoothtrj= smoothtrj
    with open('trj.pkl','w') as f:
        pickle.dump(smoothtrj,f)
    #a.animate2("./img/%02d%02d%02d%02d%02db" % (t.month, t.day, t.hour, t.minute, t.second))

    a.saveeps2("./img/%02d%02d%02d%02d%02db" % (t.month, t.day, t.hour, t.minute, t.second))
    #a.saveeps("%2d%2d%2d%2d%2d" % (t.month, t.day, t.hour, t.minute, t.second))


    # j=0
    # for t in trj:
    #     plt.plot(t[:, 0], t[:, 1], ':', color=(1 - j * 1.0 / len(trj), j * 1.0 / len(trj), j * 1.0 / len(trj)))
    #     j = j + 1

    # print len(trj[0][:,0])-1
    # for i in range(len(trj[0][:,0])-1):
    #     j=0
    #     for t in trj:
    #         plt.plot(t[i:i+2, 0], t[i:i+2, 1], ':', color=(1 - j * 1.0 / len(trj), j * 1.0 / len(trj), j * 1.0 / len(trj)))
    #         j = j + 1
        #plt.pause(0.05)
    #str = raw_input("Enter your input: ")
    #a.spantree()
    # plt.pause(100)
    # plt.show()
